import time
import cv2 as cv
import numpy as np
import sys
import torch

SNOWFLAKE, SHAKE, FROZEN, COLORCAST = 2**0, 2**1, 2**2, 2**3
NORMAL = 0


class AnomalyDetector():
    # 输入源图像
    __originImg = None
    # 池化降采样图像
    __pooledImg = None
    # 池化后灰度图像
    __grayaftpolImg = None
    # 检测采样原图像
    __sampleOriginImg = None
    # 灰度源图像
    __grayImg = None
    # 计数器
    __counter, __frozenSteps = 0, 0
    """
    异常类型
    采用逐位编码，输出范围为 (0 ~ 2 ^ 4 -1 )
    """
    __AnomalyType = 0

    __Kx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
    __Ky = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])
    __K45 = np.array([[-1, -1, 0], [-1, 0, 1], [0, 1, 1]])

    def __init__(self) -> None:
        pass

    def execute(self, img, interval=10) -> int:
        """
        execute [summary]
        视频异常检测器唯一公开接口

        Args:
            img ([cv.Mat]): 输入源图像
            interval (int, optional): 间隔采样周期. Defaults to 10.

        Returns:
            [int]: 检测到的异常类型
        """
        self.__AnomalyType = 0
        self.__pooledImg = self.__capMaxpool(img)
        self.__originImg = img
        self.__grayImg = cv.cvtColor(self.__originImg, cv.COLOR_BGR2GRAY)
        self.__grayaftpolImg = cv.cvtColor(self.__pooledImg, cv.COLOR_BGR2GRAY)

        self.__AnomalyType += self.__capColorcast(0.5) + self.__capSnowflake()
        # 周期间隔检测
        self.__counter += 1
        if self.__counter % 10 == 0:
            self.__sampleOriginImg = self.__originImg
            if self.__counter > interval:
                self.__counter = interval
                self.__AnomalyType += self.__capFrozen() + self.__capShake()
        else:
            time.sleep(0.1)

        return self.__AnomalyType

    def __capMaxpool(self, img):
        """
        __capMaxpool [summary]
        最大池化采样，目的是为了降采样提高检测速度

        Args:
            img ([cv.Mat]): 输入图像

        Returns:
            int: [cv.Mat] 降采样图像
        """
        maxp = torch.nn.MaxPool2d(kernel_size=2)
        img = torch.FloatTensor(img).permute(2, 1, 0)
        maxpedImg = maxp(img)
        maxpedImg = maxpedImg.permute(2, 1, 0).detach().numpy()
        return maxpedImg

    def __capFrozen(self, hist_thre=0.02, frame_thre=2) -> int:
        """
        __capFrozen [summary]
        画面冻结检测

        Args:
            hist_thre (float, optional): 相似度阈值. Defaults to 0.02.
            frame_thre (int, optional): 冻结帧数阈值. Defaults to 2.

        Returns:
            int: [异常类型:画面冻结]
        """
        # 转化为灰度图
        Mat1 = self.__grayImg
        Mat2 = cv.cvtColor(self.__sampleOriginImg, cv.COLOR_BGR2GRAY)
        # 计算直方图
        Mat1_hist = cv.calcHist([Mat1], [0], None, [256], [0, 256])
        Mat2_hist = cv.calcHist([Mat2], [0], None, [256], [0, 256])
        # 比较直方图
        result = cv.compareHist(Mat1_hist, Mat2_hist, cv.HISTCMP_BHATTACHARYYA)
        print("frozen_result", result)
        # 相似度越高result越小
        if result < hist_thre:
            self.__frozenSteps += 1
        else:
            self.__frozenSteps = self.__frozenSteps - 1 if self.__frozenSteps > 0 else 0
        if self.__frozenSteps > frame_thre:
            self.__frozenSteps = 0
            # 返回画面冻结信号
            return COLORCAST
        return NORMAL

    def __capColorcast(self, k_thre=0.5) -> int:
        """
        __capColorcast [summary]
        偏色检测

        Args:
            k_thre ([float]): 偏度阈值. Defaults to 0.5

        Returns:
            [int]: [异常类型:画面偏色]
        """

        """
        Lab是由一个亮度通道（channel）和两个颜色通道组成的。在Lab颜色空间中，每个颜色用L、a、b三个数字表示，各个分量的含义是这样的:
        L*代表亮度
        a*表示红色减绿色
        b*表示绿色减蓝色
        """
        m, n, _ = self.__originImg.shape
        img_lab = cv.cvtColor(self.__originImg, cv.COLOR_BGR2LAB)
        _, a, b = cv.split(img_lab)
        d_a, d_b, m_a, m_b = 0, 0, 0, 0
        for i in range(m):
            for j in range(n):
                d_a = d_a + a[i][j]
                d_b = d_b + b[i][j]
        d_a, d_b = (d_a / (m * n)) - 128, (d_b / (n * m)) - 128

        for i in range(m):
            for j in range(n):
                m_a = np.abs(a[i][j] - d_a - 128) + m_a
                m_b = np.abs(b[i][j] - d_b - 128) + m_b

        m_a, m_b = m_a / (m * n), m_b / (m * n)
        D = np.sqrt((np.square(d_a) + np.square(d_b)))
        M = np.sqrt((np.square(m_a) + np.square(m_b)))
        k = D / M
        # k 越大 偏色越高
        if k > k_thre:
            return COLORCAST
        return NORMAL

    def __capSnowflake(self) -> int:
        """
        __capSnowflake [summary]
        雪花检测

        Returns:
            [int]: [异常类型 画面雪花]
        """
        noisyp = 0
        img_x = cv.cvtColor(cv.filter2D(
            self.__pooledImg, -1, self.__Kx), cv.COLOR_BGR2GRAY)
        img_y = cv.cvtColor(cv.filter2D(
            self.__pooledImg, -1, self.__Ky), cv.COLOR_BGR2GRAY)
        img_45 = cv.cvtColor(cv.filter2D(
            self.__pooledImg, -1, self.__K45), cv.COLOR_BGR2GRAY)

        img_med = cv.cvtColor(cv.medianBlur(
            self.__pooledImg, 3), cv.COLOR_BGR2GRAY)

        for i in range(self.__grayaftpolImg.shape[0]):
            for j in range(self.__grayaftpolImg.shape[1]):
                minpixel = min(img_x[i][j], img_y[i][j], img_45[i][j])
                if minpixel > 0.1 and abs(self.__grayaftpolImg[i][j] - img_med[i][j]) > 10:
                    noisyp = noisyp + 1

        result = noisyp * 100 / \
            (self.__grayImg.shape[0] * self.__grayImg.shape[1])
        if result > 1:
            return SNOWFLAKE
        return NORMAL

    # 画面抖动检测
    # 图像灰度投影算法

    # 计算行列投影
    def __calcolrow(self):
        # 转化为灰度图
        img = cv.cvtColor(self.__originImg, cv.COLOR_BGR2GRAY)
        # 计算行的投影
        Rsum = 0
        rsum_list = []
        for i in range(img.shape[0]):
            rsum = 0
            for j in range(img.shape[1]):
                rsum = img[i][j] + rsum
            Rsum = Rsum + rsum
            rsum_list.append(rsum)
        meanR = Rsum / img.shape[0]
        # 计算列的投影
        Csum = 0
        csum_list = []
        for j in range(img.shape[1]):
            csum = 0
            for i in range(img.shape[0]):
                csum = csum + img[i][j]
            Csum = Csum + csum
            csum_list.append(csum)
        meanC = Csum / img.shape[1]
        return meanR, meanC, rsum_list, csum_list

    # 计算互相关最小值
    # m :搜索范围

    def __calmin(list1, list2, m=50):
        idx = 0
        minv = sys.maxsize
        for i in range(1, 2*m+1):
            sum_diff = 0
            for j in range((len(list1)-2*m)):
                sum_diff = pow(list1[i+j-1] - list2[j+m], 2) + sum_diff
            if (sum_diff < minv).any():
                minv = sum_diff
                idx = i
        return (m + 1 - idx)

    def __capShake(self) -> int:
        """
        __capShake [summary]
        抖动检测

        Returns:
            [int]: [异常类型 画面抖动]
        """
        samplePooledImg = self.__capMaxpool(self.__sampleOriginImg)
        new_mean_r, new_mean_c, newr_list, newc_list = self.__calcolrow(
            self.__pooledImg)
        last_mean_r, last_mean_c, lastr_list, lastc_list = self.__calcolrow(
            samplePooledImg)
        for r in range(len(newr_list)):
            newr_list[r] = newr_list[r] - new_mean_r
            lastr_list[r] = lastr_list[r] - last_mean_r
        for c in range(len(newc_list)):
            newc_list[c] = newc_list[c] - new_mean_c
            lastc_list[c] = lastc_list[c] - last_mean_c
        r_l = self.__calmin(newr_list, lastr_list)
        c_l = self.__calmin(newc_list, lastc_list)
        print("r_l: {}; c_l : {}".format(r_l, c_l))
        if abs(r_l) > 10 or abs(c_l) > 10:
            return SHAKE
        else:
            return NORMAL
